from torch.utils.data import Dataset

from utils.DataUtils import get_patch_with_mask
from utils.LogUtil import my_logger


class CustomDataset(Dataset):
    def __init__(self, img_list, label_list, dataset_type="train", ):
        self.img_list = img_list
        self.label_list = label_list
        self.dataset_type = dataset_type
        my_logger.info("The " + self.dataset_type + " data's size is " + str(len(self.img_list)))

    def __len__(self):
        return len(self.img_list)

    def __getitem__(self, idx):
        img, label = self.img_list[idx], self.label_list[idx]

        # img = torch.from_numpy(img)
        # label = torch.from_numpy(label).long()
        # label = torch.nn.functional.one_hot(label)  # The shape is [W,H,L,C]

        return img, label


class CustomLazyDataset(Dataset):
    def __init__(self, sample_list, patch_shape=(32, 64, 64), crop_method=0, dataset_type="train"):
        """
        The patch_shape only work if if_random_crop = True
        :param sample_list:
        :param patch_shape:
        :param if_random_crop:
        """
        self.image_list = []
        self.label_list = []
        self.sample_list = sample_list
        self.dataset_type = dataset_type
        self.patch_shape = patch_shape
        self.crop_method = crop_method
        my_logger.info("The " + self.dataset_type + " data's size is " + str(len(self.sample_list)))

    def __len__(self):
        return len(self.sample_list)

    def __getitem__(self, idx):
        if len(self.image_list) == 0:
            sample_path = self.sample_list.pop()
            tmp_image_list, tmp_label_list = get_patch_with_mask(sample_path,
                                                                 self.patch_shape,
                                                                 self.crop_method)
            self.image_list += tmp_image_list
            self.label_list += tmp_label_list
        return self.image_list.pop(), self.label_list.pop()
        # my_logger.info(str(len(tmp_data_list)) + " patches has been added to list")
